
import numpy as np
import cv2
import glob

#棋盘格点阵
ROW=7       #行
COL=9      #列
#终止条件
criteria=(cv2.TERM_CRITERIA_EPS +cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
#准备对象点，like(0,0,0),(1,0,0),(2,0,0),...,(6,5,0)
#这其实就是棋盘格点的空间坐标点，令其Z=0，而相机中心是变动的。
#实际的棋盘方格尺寸可能不是1x1，这里我们归一化到1x1。
objp=np.zeros((ROW*COL,3), np.float32)
objp[:,:2]=np.mgrid[0:ROW, 0:COL].T.reshape(-1,2)
objpoints=[]    #现实空间中的3D点
impoints=[]     #像平面上的2D点
images=glob.glob(r'f:\chess\IMG*.jpg')
width=0
height=0
for fname in images:
    img=cv2.imread(fname)
    gray=cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)   #变成灰度图
    #寻找棋盘的角点, 只匹配其中的7x6的角点阵。
    ret,corners=cv2.findChessboardCorners(gray,(ROW,COL),None)      #cv2.CALIB_CB_NORMALIZE_IMAGE+cv2.CALIB_CB_ADAPTIVE_THRESH+cv2.CALIB_CB_FILTER_QUADS
    if ret==True:
        objpoints.append(objp)      #添加7*6=42个点的数据，以备后面使用
        #对找到的角点进行细化，提高精度。
        corners2=cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
        impoints.append(corners2)   #添加这些角点到像点数组
        #绘制这些角点
        img=cv2.drawChessboardCorners(img,(ROW,COL),corners2,ret)
        while True:
            cv2.imshow('img',img)
            if cv2.waitKey() & 0xff == ord("q"):
                break
    cv2.destroyAllWindows()
#进行校准
#返回值：
# ret表示校准是否成功，
# mtx为得到的相机矩阵，
# dist为畸变系数
# rvecs为每一个棋盘图像的旋转向量
# tvecs为每一个棋盘图像的平移向量
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, impoints, gray.shape[::-1],None,None)

##优化
img = cv2.imread(images[0])
height,width = img.shape[:2]
#newcameramtx, roi=cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1,(w,h))

##对图像进行校正
#img_dst = cv2.undistort(img, mtx, None, None, newcameramtx)
img_dst = cv2.undistort(img, mtx, None, None, None)
##生成校正后的图像
#x,y,w,h = roi
#img_dst = img_dst[y:y+h, x:x+w]
cv2.imwrite(r'chess_pic\calibresult.png',img_dst)
print(mtx)
print(dist)
#print(newcameramtx)